/*
Copyright (C) 2011 by the PSVN Research Group, University of Alberta
*/

#ifndef _HIDA_HPP
#define _HIDA_HPP

#include <vector>
#include <string>
#include "abst_desc.hpp"
#include "hash_map.hpp"
#include "reservoir.hpp"

class FringeEntry
{
public:
    int64_t hash;
    int g;
    FringeEntry() {};
    FringeEntry(int64_t hash, int g) : hash(hash), g(g) { }
};

struct FringeCacheData
{
    int h;
    int exact;
    int g;
    int search_id;
    void* state;
    void* parent;
    FringeCacheData() : h(-1), exact(0), g(0), search_id(-1),
                        state(0), parent(0) { };
    bool IsValid() const { return search_id != -1; }
    void Invalidate() { search_id = -1; }
};

struct HidaCacheData
{
    int h;
    int exact;
    HidaCacheData() : h(-1), exact(0) { }
    bool IsValid() const { return h != -1; }
    void Invalidate() { h = -1; }
};

template<class AH>
class HIDA
{
public:
    HIDA(bool use_fringe, const AH& abst);
    ~HIDA();

    int Solve(const void *orig_state, std::vector<const void*>* solution);
    void ClearCache();

    void NumExpanded(std::vector<size_t>& expanded) const;
    void NumWasted(std::vector<size_t>& wasted) const;
    void NumStored(std::vector<size_t>& stored) const;
    void NumSearches(std::vector<size_t>& searches) const;
    void Memory(std::vector<size_t>& memory) const;

private:
    static const int MAX_DEPTH = 128;

    class Level
    {
    public:
        HashMap<HidaCacheData> hida_data;
        HashMap<FringeCacheData> fringe_data;
        std::vector<FringeEntry> fringe[2];
        std::vector<int64_t> expanded;
        std::vector<void*> state;
        std::vector<int64_t> history;
        int search_id;
        uint64_t nodes_expanded;
        uint64_t nodes_expanded_this_search;
        uint64_t wasted_expansions;
        uint64_t searches;

        Level()
            : state(HIDA::MAX_DEPTH)
            , history(HIDA::MAX_DEPTH)      
            , search_id(-1)
        { };
    };

    const AH& m_ah;
    Level* m_levels;
    Reservoir m_memory;
    bool m_use_fringe;
    bool m_use_tree_backup;
    bool m_use_Pg_caching;

    std::vector<const void*>* m_solution;

    void* NewState()
    { return m_memory.Get(m_ah.StateSize()); }

    int SearchID(const int level) const 
    { return m_levels[level].search_id; }

    //------------------------------------------------------
    // HIDA cache
    //------------------------------------------------------
    bool GetHidaCache(const int level, const int64_t hash,
                      HidaCacheData& data) const
    { return m_levels[level].hida_data.Get(hash, data); }

    HidaCacheData* GetHidaCachePointer(const int level, const int64_t hash)
    { return m_levels[level].hida_data.GetPointer(hash); }

    void StoreInHidaCache(const int level, const int64_t hash, 
                          const HidaCacheData& data)
    { m_levels[level].hida_data.Add(hash, data); }

    //------------------------------------------------------
    // Fringe cache
    //------------------------------------------------------
    bool GetFringeCache(int level, int64_t hash, FringeCacheData& data) const
    { return m_levels[level].fringe_data.Get(hash, data); }

    FringeCacheData* GetFringeCachePointer(int level, int64_t hash) 
    { return m_levels[level].fringe_data.GetPointer(hash); }

    void StoreInFringeCache(int level, int64_t hash,const FringeCacheData& data)
    { m_levels[level].fringe_data.Add(hash, data); }


    void StoreSolution(const int level, const int depth) const;
    bool StateInHistory(const int level, int depth, const int64_t hash) const;

    int Heuristic(const int level, const void* state);
    bool DFS(const int level, const void* state, const bool check_cycles,
             const int depth, const int g, const int bound, 
             int* new_bound, int* new_estimate);
    int IDA(const int level, const void* state);
    int Search(const int level, const void* state);

    bool FringeSearch(const int level, const void* state, int* bound);
    bool BackupF(const int level, const void* state, 
                 const int depth, const int g, const int bound, int* need);

    /* Non-copyable */
    HIDA(const HIDA& other);
    void operator=(const HIDA& other);
};

template<class AH>
HIDA<AH>::HIDA(bool use_fringe, const AH& abst)
    : m_ah(abst)
    , m_levels(new Level[abst.NumLevels()])
    , m_memory(16)  // allocate in 16MB chunks
    , m_use_fringe(use_fringe)
    , m_use_tree_backup(false)
    , m_use_Pg_caching(true)
{
    ClearCache();
}

template<class AH>
HIDA<AH>::~HIDA()
{
    delete [] m_levels;
}

template<class AH>
void HIDA<AH>::ClearCache()
{
    for (size_t i = 0; i < m_ah.NumLevels(); ++i) {
        m_levels[i].hida_data.Clear();
        m_levels[i].fringe_data.Clear();
    }
    m_memory.Reset();
    for (size_t i = 0; i < m_ah.NumLevels(); ++i) {
        for (int j = 0; j < MAX_DEPTH; ++j)
            m_levels[i].state[j] = NewState();
    }
}

template<class AH>
void HIDA<AH>::StoreSolution(const int level, const int depth) const
{
    for (int i = 0; i <= depth; ++i)
        m_solution->push_back(m_levels[level].state[i]);
}

template<class AH>
bool HIDA<AH>::StateInHistory(const int level, int depth, 
                              const int64_t hash) const
{
    for (--depth; depth >= 0; --depth)
        if (hash == m_levels[level].history[depth])
            return true;
    return false;
}

template<class AH>
int HIDA<AH>::Heuristic(const int level, const void* state)
{
    if (level == int(m_ah.NumLevels() - 1))
        return 0;
    void* abstracted_state = m_levels[level + 1].state[0];
    m_ah.AbstractToNext(level, state, abstracted_state);
    int64_t abst_hash = m_ah.Hash(abstracted_state);
    if (m_use_fringe) {
        FringeCacheData data;
        if (GetFringeCache(level + 1, abst_hash, data) && data.exact)
            return data.h;
    } else {
        HidaCacheData data;
        if (GetHidaCache(level + 1, abst_hash, data) && data.exact)
            return data.h;
    }
    return Search(level + 1, abstracted_state);
}

template<class AH>
bool HIDA<AH>::FringeSearch(const int level, const void* state, int* bound)
{
    ++m_levels[level].search_id;

    int64_t hash = m_ah.Hash(state);
    FringeCacheData rootdata;
    if (!GetFringeCache(level, hash, rootdata)) {
        rootdata.state = NewState();
        m_ah.Copy(rootdata.state, state);
    }
    assert(!rootdata.exact);
    rootdata.g = 0;
    rootdata.search_id = SearchID(level);
    rootdata.parent = NULL;
    if (rootdata.h == -1) {
        rootdata.h = Heuristic(level, state);
    }
    StoreInFringeCache(level, hash, rootdata);
    *bound = rootdata.h;

    std::vector<FringeEntry>* now = &m_levels[level].fringe[0];
    std::vector<FringeEntry>* later = &m_levels[level].fringe[1];
    std::vector<int64_t>* expanded = &m_levels[level].expanded;

    now->clear();
    later->clear();
    expanded->clear();

    now->push_back(FringeEntry(hash, 0));
    bool solved = false;
    while (!solved && !now->empty()) {
        int next_bound = INT_MAX;
        while (!now->empty()) {
            FringeCacheData data;
            const int64_t current_hash = now->back().hash;
            const int64_t fringe_g = now->back().g;
            now->pop_back();
            if (!GetFringeCache(level, current_hash, data))
                assert(false);

            if (fringe_g > data.g) { 
                // skip this one... we already saw a better way to get here
                m_levels[level].wasted_expansions++;
                continue;
            }
            if (data.g > *bound) {
                next_bound = std::min(next_bound, data.g);
                later->push_back(FringeEntry(current_hash, data.g));
                continue;
            }
            if (data.h == -1) {   // compute heuristic if necessary 
                data.h = std::max(m_ah.Epsilon(level, data.state), 
                                  Heuristic(level, data.state));
                data.exact = (data.h == 0);
                StoreInFringeCache(level, current_hash, data);
            }
            if (data.g + data.h > *bound) { // too big, do it later
                next_bound = std::min(next_bound, data.g + data.h);
                later->push_back(FringeEntry(current_hash, data.g));
                continue;
            } 
            if ((data.exact && (data.g + data.h == *bound)) 
                || m_ah.IsGoal(level, data.state)) {
                solved = true;
                break;
            }
            m_levels[level].nodes_expanded++;
            m_levels[level].nodes_expanded_this_search++;
            expanded->push_back(current_hash);
            void* child = m_levels[level].state[0];
            typename AH::FwdReverseChildGenerator 
                generator(m_ah, level, data.state);
            while (generator.Generate(child)) {
                int child_g = data.g + generator.Cost();
                int64_t child_hash = m_ah.Hash(child);
                if (FringeCacheData* child_ptr = GetFringeCachePointer(level, child_hash)) {
                    if (child_ptr->search_id != SearchID(level)) {
                        // saw child in a previous search; we
                        // can reuse the h value, but nothing else.
                        child_ptr->g = child_g;
                        child_ptr->parent = data.state;
                        child_ptr->search_id = SearchID(level);
                        now->push_back(FringeEntry(child_hash, child_g));
                    } else {
                        // saw child in this search already, update
                        // to shorter g if necessary
                        if (child_g < child_ptr->g) {
                            child_ptr->g = child_g;
                            child_ptr->parent = data.state;
                            now->push_back(FringeEntry(child_hash, child_g));
                        }
                    }
                } else {
                    // have never seen this state before
                    FringeCacheData child_data;
                    child_data.h = -1; // no heuristic info yet
                    child_data.g = child_g;
                    child_data.parent = data.state;
                    child_data.search_id = SearchID(level);
                    child_data.state = NewState();
                    m_ah.Copy(child_data.state, child);
                    StoreInFringeCache(level, child_hash, child_data);
                    now->push_back(FringeEntry(child_hash, child_g));
                }
            }
        }
        if (!solved)
            *bound = next_bound;
        std::swap(now, later);
    }
    assert(solved);
    if (m_use_tree_backup) {
        // backup f-values using a simple dfs
        int new_bound;
        bool dfssolved = BackupF(level, rootdata.state, 0, 0, 
                                 *bound, &new_bound);
        assert(dfssolved == solved);
    } else {
        // do P-g caching for expanded nodes
        for (size_t i = 0; i < expanded->size(); ++i) {
            FringeCacheData* ptr = GetFringeCachePointer(level, (*expanded)[i]);
            assert(ptr);
            assert(ptr->search_id == SearchID(level));
            ptr->h = *bound - ptr->g;
        }
        // mark nodes on solution as exact
        if (solved) {
            const void* state = rootdata.state;
            while (state) {
                int64_t hash = m_ah.Hash(state);
                FringeCacheData* ptr = GetFringeCachePointer(level, hash);
                assert(ptr);
                assert(ptr->search_id == SearchID(level));
                ptr->h = *bound - ptr->g;
                ptr->exact = true;
                state = ptr->parent;
            }
        }
    }
    return solved;
}

template<class AH>
bool HIDA<AH>::BackupF(const int level, const void* state, 
                   const int depth, const int g, const int bound, int* need)
{
    assert(depth < MAX_DEPTH);

    FringeCacheData data;
    int64_t hash = m_ah.Hash(state);
    bool cached = GetFringeCache(level, hash, data);
    if (!cached || data.h == -1) {
        *need = bound;
        return false;
    }
    if (bound > data.h) {
        // P-g caching
        data.h = bound;
        StoreInFringeCache(level, hash, data);
    }
    if (data.search_id == SearchID(level) && data.g < g) {
        *need = data.h;
        return false;
    }
    if (data.h > bound) {
        *need = data.h;
        return false;
    } else if (data.h == bound && data.exact) {
        *need = data.h;
        return true;
    }

    *need = INT_MAX;
    void* child = m_levels[level].state[depth + 1];
    typename AH::FwdChildGenerator generator(m_ah, level, state);
    while (generator.Generate(child)) {
        const int cost = generator.Cost();
        if (cost > bound)
            *need = std::min(*need, cost);
        else {
            int child_need;
            if (BackupF(level, child, depth + 1, g + cost, 
                        bound - cost, &child_need)) {
                *need = bound;
                data.exact = true;
                break;
            }
            *need = std::min(*need, cost + child_need);
        }
    }
    data.h = *need;
    StoreInFringeCache(level, hash, data);
    return data.exact;
}

template<class AH>
bool HIDA<AH>::DFS(const int level, const void* state, const bool check_cycles,
               const int depth, const int g, const int bound, 
               int* new_bound, int* new_estimate)
{
    assert(depth < MAX_DEPTH);

    m_levels[level].history[depth] = m_ah.Hash(state);
    const int64_t& hash = m_levels[level].history[depth];

    if (HidaCacheData* cache_ptr = GetHidaCachePointer(level, hash)) {
        // check for cut-off
        if (cache_ptr->h > bound) {
            *new_bound = *new_estimate = cache_ptr->h;
            return false;
        }
        // check for solution
        else if (cache_ptr->h == bound && cache_ptr->exact) {
            *new_bound = *new_estimate = cache_ptr->h;
            if (level == 0) 
                StoreSolution(level, depth);
            return true;
        }
        // do P-g caching
        if (m_use_Pg_caching && (bound > cache_ptr->h)) {
            cache_ptr->h = bound;
        }
        // check for 0-cost cycles
        if (check_cycles && StateInHistory(level, depth, hash)) {
            *new_estimate = cache_ptr->h;
            *new_bound = INT_MAX;
            return false;
        }
    } else {
        HidaCacheData data;
        // have never been to this state, so compute heuristic
        data.h = std::max(bound, std::max(m_ah.Epsilon(level, state), 
                                          Heuristic(level, state)));
        data.exact = (data.h == 0) && m_ah.IsGoal(level, state);
        StoreInHidaCache(level, hash, data);

        // check for cutoff
        if (data.h > bound) {
            *new_bound = *new_estimate = data.h;
            return false;
        } else if (data.h == bound && data.exact) {
            *new_bound = *new_estimate = data.h;
            if (level == 0)
                StoreSolution(level, depth);
            return true;
        }
    }
    // expand children
    bool solved = false;
    *new_bound = *new_estimate = INT_MAX;
    m_levels[level].nodes_expanded++;
    m_levels[level].nodes_expanded_this_search++;
    void* child = m_levels[level].state[depth + 1];
    typename AH::FwdChildGenerator generator(m_ah, level, state);
    while (generator.Generate(child)) {
        const int cost = generator.Cost();
        if (cost > bound) {
            *new_bound = std::min(*new_bound, cost);
            *new_estimate = std::min(*new_estimate, cost);
        }
        else {
            int child_bound, child_estimate;
            if (DFS(level, child, cost == 0, depth + 1, g + cost, 
                    bound - cost, &child_bound, &child_estimate)) 
            {
                *new_bound = *new_estimate = bound;
                solved = true;
                break;
            }
            if (child_bound != INT_MAX)
                *new_bound = std::min(*new_bound, cost + child_bound);
            if (child_estimate != INT_MAX)
                *new_estimate = std::min(*new_estimate, cost + child_estimate);
        }
    }
    // store updated estimate
    HidaCacheData* ptr = GetHidaCachePointer(level, hash);
    assert(solved || (*new_bound > ptr->h));
    ptr->h = *new_estimate;
    ptr->exact = solved;
    return solved;
}

template<class AH>
int HIDA<AH>::IDA(const int level, const void* state)
{
    int bound = Heuristic(level, state);
    m_levels[level].search_id++;
    uint64_t wasted = 0;
    while (bound != INT_MAX) {
        int new_bound, new_estimate;
        m_levels[level].nodes_expanded_this_search = 0;
        if (DFS(level, state, false, 0, 0, bound, &new_bound, &new_estimate))
            break;
        wasted += m_levels[level].nodes_expanded_this_search;
        bound = new_bound;
    }
    m_levels[level].wasted_expansions += wasted;
    return bound;
}

template<class AH>
int HIDA<AH>::Search(const int level, const void* state)
{
    m_levels[level].searches++;
    if (m_use_fringe) {
        int bound;
        FringeSearch(level, state, &bound);
        return bound;
    }
    return IDA(level, state);
}

template<class AH>
int HIDA<AH>::Solve(const void* orig_state, std::vector<const void*>* solution) 
{    
    for (size_t i = 0; i < m_ah.NumLevels(); ++i) {
        m_levels[i].nodes_expanded = 0;
        m_levels[i].wasted_expansions = 0;
        m_levels[i].searches = 0;
    }
    m_ah.Copy(m_levels[0].state[0], orig_state);
    m_solution = solution;
    m_solution->clear();
    return IDA(0, m_levels[0].state[0]);
}

template<class AH>
void HIDA<AH>::NumExpanded(std::vector<size_t>& expanded) const
{
    expanded.resize(m_ah.NumLevels());
    for (size_t i = 0; i < m_ah.NumLevels(); ++i)
        expanded[i] = m_levels[i].nodes_expanded;
}

template<class AH>
void HIDA<AH>::NumWasted(std::vector<size_t>& wasted) const
{
    wasted.resize(m_ah.NumLevels());
    for (size_t i = 0; i < m_ah.NumLevels(); ++i)
        wasted[i] = m_levels[i].wasted_expansions;
}

template<class AH>
void HIDA<AH>::NumStored(std::vector<size_t>& stored) const
{
    stored.resize(m_ah.NumLevels());
    for (size_t i = 0; i < m_ah.NumLevels(); ++i)
        stored[i] = (i == 0 || !m_use_fringe) 
            ? m_levels[i].hida_data.NumStored()
            : m_levels[i].fringe_data.NumStored();
}

template<class AH>
void HIDA<AH>::NumSearches(std::vector<size_t>& searches) const
{
    searches.resize(m_ah.NumLevels());
    for (size_t i = 0; i < m_ah.NumLevels(); ++i)
        searches[i] = m_levels[i].searches;
}

template<class AH>
void HIDA<AH>::Memory(std::vector<size_t>& mem) const
{
    mem.resize(m_ah.NumLevels());
    for (size_t i = 0; i < m_ah.NumLevels(); ++i) {
        mem[i] = (i == 0 || !m_use_fringe) 
            ? m_levels[i].hida_data.Memory() 
              + m_levels[i].hida_data.NumStored() * m_ah.StateSize()
            : m_levels[i].fringe_data.Memory() 
              + m_levels[i].fringe_data.NumStored() * m_ah.StateSize();
    }
}

#endif // _HIDA_HPP
